"""
矿山安全评估模型推理Handler

Date: 2025-10-19
"""

import joblib
import numpy as np
import pandas as pd
from pathlib import Path
from typing import Union, List
import warnings

warnings.filterwarnings("ignore")


class MineEvaluationHandler:
    """
    矿山安全评估模型推理处理器
    
    该类封装了模型加载和推理的完整流程，提供标准化的预测接口。
    
    Attributes:
        model: 训练好的MLP分类器
        label_encoder: 标签编码器，用于将预测结果转换为原始标签
        feature_scaler: 特征标准化器，用于对输入数据进行标准化
        selected_features: 模型使用的特征列表
        feature_names: 按顺序排列的特征名称列表
    """
    
    # 定义所有原始特征（用于scaler标准化，必须与训练时一致）
    ALL_FEATURES = [
        '工人的安全意识与操作能力',
        '管理人员的管理能力',
        '技术人员的技术素质',
        '特殊岗位人员持证率',
        '采掘设备完好率',
        '支护设备完好率',
        '采掘机械化水平',
        '通风构筑物完好率',
        '矿井通风设备完好率',
        '瓦斯监测设备完好率',
        '瓦斯抽放设备完好率',
        '隔爆防爆设备完好率',
        '火灾监测设备完好率',
        '防火及消防设施完好率',
        '探防水设备完好率',
        '井下排水设备完好率',
        '运输设备完好率',
        '运输机械化水平',
        '矿井提升设备完好率',
        '供电系统设备完好率',
        '供电系统保护装置完好率',
        '煤层自燃发火倾向性',
        '平均瓦斯涌出量',
        '矿井正常涌水量',
        '风量供需比',
        '顶板管理难易程度',
        '安全管理制度及落实率',
        '作业人员培训及考核',
        '安全检查技术人员配备率',
        '安全投入指数',
        '应急响应机制完善率',
        '安全检查整改落实率'
    ]
    
    # 定义用户输入的15个关键特征名称和顺序
    FEATURE_NAMES = [
        '工人的安全意识与操作能力',
        '管理人员的管理能力',
        '采掘设备完好率',
        '矿井通风设备完好率',
        '瓦斯监测设备完好率',
        '防火及消防设施完好率',
        '井下排水设备完好率',
        '矿井提升设备完好率',
        '供电系统设备完好率',
        '风量供需比',
        '顶板管理难易程度',
        '煤层自燃发火倾向性',
        '安全管理制度及落实率',
        '应急响应机制完善率',
        '安全检查整改落实率'
    ]
    
    def __init__(self, model_dir: str = "./saved_models", 
                 timestamp: str = "20251019_150609"):
        """
        初始化模型处理器，加载预训练的模型和相关组件
        
        Args:
            model_dir: 模型文件所在目录路径
            timestamp: 模型文件的时间戳标识
            
        Raises:
            FileNotFoundError: 当模型文件不存在时抛出
            Exception: 当模型加载失败时抛出
        """
        self.model_dir = Path(model_dir)
        self.timestamp = timestamp
        
        # 构建模型文件路径
        model_path = self.model_dir / f"mlp_model_{timestamp}.joblib"
        encoder_path = self.model_dir / f"label_encoder_{timestamp}.joblib"
        scaler_path = self.model_dir / f"feature_scaler_{timestamp}.joblib"
        features_path = self.model_dir / f"selected_features_{timestamp}.joblib"
        
        # 验证文件存在性
        self._validate_files([model_path, encoder_path, scaler_path, features_path])
        
        # 加载模型组件
        try:
            self.model = joblib.load(model_path)
            self.label_encoder = joblib.load(encoder_path)
            self.feature_scaler = joblib.load(scaler_path)
            self.selected_features = joblib.load(features_path)
            
            print(f"✓ 模型加载成功！")
            print(f"  - 模型时间戳: {timestamp}")
            print(f"  - 原始特征数量: {len(self.ALL_FEATURES)}")
            print(f"  - 输入特征数量: {len(self.FEATURE_NAMES)}")
            print(f"  - 模型使用特征数量: {len(self.selected_features)}")
            print(f"  - 类别数量: {len(self.label_encoder.classes_)}")
            print(f"  - 可预测类别: {list(self.label_encoder.classes_)}")
            
        except Exception as e:
            raise Exception(f"模型加载失败: {str(e)}")
    
    def _validate_files(self, file_paths: List[Path]) -> None:
        """
        验证模型文件是否存在
        
        Args:
            file_paths: 待验证的文件路径列表
            
        Raises:
            FileNotFoundError: 当任一文件不存在时抛出
        """
        for path in file_paths:
            if not path.exists():
                raise FileNotFoundError(f"模型文件不存在: {path}")
    
    def _validate_dataframe(self, df: pd.DataFrame) -> None:
        """
        验证输入DataFrame的格式和内容
        
        Args:
            df: 待验证的DataFrame
            
        Raises:
            ValueError: 当DataFrame格式不符合要求时抛出
        """
        if not isinstance(df, pd.DataFrame):
            raise ValueError(f"输入必须是pandas DataFrame，当前类型: {type(df)}")
        
        if df.empty:
            raise ValueError("输入DataFrame不能为空")
        
        missing_cols = set(self.FEATURE_NAMES) - set(df.columns)
        if missing_cols:
            raise ValueError(f"DataFrame缺少必需的列: {missing_cols}")
        
        # 检查数据类型
        for col in self.FEATURE_NAMES:
            if not pd.api.types.is_numeric_dtype(df[col]):
                raise ValueError(f"列 '{col}' 必须是数值类型")
    
    def predict(self, df: pd.DataFrame) -> np.ndarray:
        """
        对输入的DataFrame进行批量预测
        
        Args:
            df: 包含所有必需特征的DataFrame，列名必须与FEATURE_NAMES匹配
            
        Returns:
            预测结果数组，包含每个样本的预测类别（原始标签）
            
        Raises:
            ValueError: 当输入格式不正确时抛出
            
        Example:
            >>> handler = MineEvaluationHandler()
            >>> df = pd.DataFrame({...})  # 包含15个特征的DataFrame
            >>> predictions = handler.predict(df)
            >>> print(predictions)
            ['I类' 'II类' ...]
        """
        # 验证输入
        self._validate_dataframe(df)
        
        # 按照标准顺序提取用户输入的15个特征
        X_input = df[self.FEATURE_NAMES].copy()
        
        # 创建包含所有32个原始特征的DataFrame，缺失的特征填充为0
        X_full = pd.DataFrame(0, index=X_input.index, columns=self.ALL_FEATURES)
        
        # 将用户输入的15个特征填充到对应位置
        for col in self.FEATURE_NAMES:
            X_full[col] = X_input[col]
        
        # 使用完整的32个特征进行标准化
        X_scaled = self.feature_scaler.transform(X_full)
        
        # 将标准化后的数据转换为DataFrame
        X_scaled_df = pd.DataFrame(X_scaled, columns=self.ALL_FEATURES, index=X_input.index)
        
        # 提取模型需要的特征（selected_features）
        X_model = X_scaled_df[self.selected_features]
        
        # 进行预测（编码后的标签）
        y_pred_encoded = self.model.predict(X_model)
        
        # 解码为原始标签
        y_pred = self.label_encoder.inverse_transform(y_pred_encoded)
        
        return y_pred
    
    def predict_one(self, features: Union[List[float], np.ndarray]) -> str:
        """
        对单个样本进行预测
        
        Args:
            features: 包含15个特征值的列表或数组（未标准化的原始值）
                     顺序必须与FEATURE_NAMES一致
            
        Returns:
            预测的类别标签（字符串）
            
        Raises:
            ValueError: 当输入特征数量不正确时抛出
            
        Example:
            >>> handler = MineEvaluationHandler()
            >>> features = [0.85, 0.90, 0.88, 0.92, 0.87, 0.89, 0.91, 0.88, 0.90,
            ...            1.2, 0.75, 0.8, 0.95, 0.92, 0.93]
            >>> prediction = handler.predict_one(features)
            >>> print(f"预测类别: {prediction}")
            预测类别: I类
        """
        # 验证输入
        if not isinstance(features, (list, np.ndarray)):
            raise ValueError(f"输入必须是列表或numpy数组，当前类型: {type(features)}")
        
        features = np.array(features).flatten()
        
        if len(features) != len(self.FEATURE_NAMES):
            raise ValueError(
                f"特征数量不匹配。期望: {len(self.FEATURE_NAMES)}, 实际: {len(features)}"
            )
        
        # 构建DataFrame（使用标准特征名称）
        df = pd.DataFrame([features], columns=self.FEATURE_NAMES)
        
        # 调用批量预测方法
        prediction = self.predict(df)
        
        return prediction[0]
    
    def predict_proba(self, df: pd.DataFrame) -> np.ndarray:
        """
        预测各类别的概率分布
        
        Args:
            df: 包含所有必需特征的DataFrame
            
        Returns:
            概率矩阵，shape为(n_samples, n_classes)
            
        Example:
            >>> handler = MineEvaluationHandler()
            >>> df = pd.DataFrame({...})
            >>> probas = handler.predict_proba(df)
            >>> print(f"各类别概率: {probas[0]}")
        """
        # 验证输入
        self._validate_dataframe(df)
        
        # 按照标准顺序提取用户输入的15个特征
        X_input = df[self.FEATURE_NAMES].copy()
        
        # 创建包含所有32个原始特征的DataFrame，缺失的特征填充为0
        X_full = pd.DataFrame(0, index=X_input.index, columns=self.ALL_FEATURES)
        
        # 将用户输入的15个特征填充到对应位置
        for col in self.FEATURE_NAMES:
            X_full[col] = X_input[col]
        
        # 使用完整的32个特征进行标准化
        X_scaled = self.feature_scaler.transform(X_full)
        
        # 将标准化后的数据转换为DataFrame
        X_scaled_df = pd.DataFrame(X_scaled, columns=self.ALL_FEATURES, index=X_input.index)
        
        # 提取模型需要的特征（selected_features）
        X_model = X_scaled_df[self.selected_features]
        
        # 预测概率
        probas = self.model.predict_proba(X_model)
        
        return probas
    
    def get_feature_names(self) -> List[str]:
        """
        获取模型使用的特征名称列表
        
        Returns:
            特征名称列表
        """
        return self.FEATURE_NAMES.copy()
    
    def get_classes(self) -> np.ndarray:
        """
        获取模型可预测的类别列表
        
        Returns:
            类别数组
        """
        return self.label_encoder.classes_.copy()
    
    def get_selected_features(self) -> List[str]:
        """
        获取模型实际使用的特征列表
        
        Returns:
            模型使用的特征名称列表
        """
        return self.selected_features.copy()
    
    def print_feature_info(self) -> None:
        """
        打印特征处理流程的详细信息
        
        该方法展示从输入特征到模型特征的完整转换过程
        """
        print("\n" + "=" * 70)
        print("特征处理流程信息")
        print("=" * 70)
        print(f"\n步骤1: 用户输入 {len(self.FEATURE_NAMES)} 个特征")
        print("-" * 70)
        for i, feat in enumerate(self.FEATURE_NAMES, 1):
            print(f"  {i:2d}. {feat}")
        
        print(f"\n步骤2: 补全为 {len(self.ALL_FEATURES)} 个原始特征（缺失的填0）")
        print("-" * 70)
        missing_features = [f for f in self.ALL_FEATURES if f not in self.FEATURE_NAMES]
        print(f"  补充的特征（共 {len(missing_features)} 个，填充为0）:")
        for i, feat in enumerate(missing_features, 1):
            print(f"  {i:2d}. {feat}")
        
        print(f"\n步骤3: 使用StandardScaler对32个特征进行标准化")
        print("-" * 70)
        
        print(f"\n步骤4: 提取模型需要的 {len(self.selected_features)} 个特征进行预测")
        print("-" * 70)
        for i, feat in enumerate(self.selected_features, 1):
            print(f"  {i:2d}. {feat}")
        
        print("\n" + "=" * 70 + "\n")



def demo_usage():
    """
    演示如何使用MineEvaluationHandler
    """
    print("\n" + "=" * 60)
    print("使用示例演示")
    print("=" * 60 + "\n")
    
    # 初始化handler
    handler = MineEvaluationHandler()
    # handler = MineEvaluationHandler(model_dir="./AHP-MIV-BP_0.93/saved_models")
    
    # 示例1：单样本预测
    print("示例1: 单样本预测!")
    print("-" * 40)
    features = [0.85, 0.90, 0.88, 0.92, 0.87, 0.89, 0.91, 0.88, 0.90,
                1.2, 0.75, 0.8, 0.95, 0.92, 0.93]
    result = handler.predict_one(features)
    print(f"输入特征: {features[:3]}... (共15个)")
    print(f"预测类别: {result}\n")
    
    # 示例2：批量预测
    print("示例2: 批量预测")
    print("-" * 40)
    df = pd.DataFrame({
        '工人的安全意识与操作能力': [0.85, 0.70,0,10],
        '管理人员的管理能力': [0.90, 0.75,0,10],
        '采掘设备完好率': [0.88, 0.80,0,10],
        '矿井通风设备完好率': [0.92, 0.78,0,10],
        '瓦斯监测设备完好率': [0.87, 0.76,0,10],
        '防火及消防设施完好率': [0.89, 0.79,0,10],
        '井下排水设备完好率': [0.91, 0.81,0,10],
        '矿井提升设备完好率': [0.88, 0.77,0,10],
        '供电系统设备完好率': [0.90, 0.80,0,10],
        '风量供需比': [1.2, 0.9,0,10],
        '顶板管理难易程度': [0.75, 0.6,0,10],
        '煤层自燃发火倾向性': [0.8, 0.65,0,10],
        '安全管理制度及落实率': [0.95, 0.82,0,10],
        '应急响应机制完善率': [0.92, 0.80,0,10],
        '安全检查整改落实率': [0.93, 0.81,0,10]
    })
    results = handler.predict(df)
    print(f"输入样本数: {len(df)}")
    print(f"预测结果: {results}\n")


if __name__ == "__main__":

    
    # 运行使用示例
    demo_usage()

